import streamlit as st
from streamlit_chatbox import ChatBox,Markdown
from openai import OpenAI
import os

# 可直接替换为api地址和key
API_BASE = os.environ.get("YI_AI_API_BASE")
API_KEY = os.environ.get("YI_AI_API_KEY")

client = OpenAI(
    api_key=API_KEY,
    base_url=API_BASE
)

# 系统扮演角色选择
system_roles = {
    "Chat Robot": "你是一个智能机器人，精通各行业的知识。可以回答用户提出的任何问题。",
    "Java to Python": "你是一个精通Java和Python的软件开发工程师，现在需要帮助一个只熟悉Java开发语言的程序员学习Python语言开发。请结合Java的知识帮助他快速入门。你需要使用中文回答所有问题。",
    "Java to C/C++": "请扮演一位精通Java、C和C++的软件开发工程师，目标是帮助一位仅熟悉Java语言的程序员快速入门C和C++。在指导过程中，请充分利用学习者已有的Java知识作为桥梁，通过对比Java与C/C++之间的异同来促进理解。具体来说，可以从基本语法、数据类型、内存管理、面向对象编程等方面入手，并提供相应的代码示例以加深印象。此外，如果可能的话，请推荐一些适合初学者的学习资源或工具，以便于他能够更加系统地掌握这两种新语言。整个过程请使用中文进行交流。",
    "Social Skills Expert": "请为一位熟悉Java但希望学习Python的程序员提供入门指导。在你的指导中，请结合他已有的Java知识来解释Python的概念，以便于他能够更快地理解和掌握Python编程语言的基础。请确保使用中文进行回答，并尽可能地将两种语言之间的相似点和不同点清晰地呈现出来，帮助他顺利过渡到Python开发。此外，如果可能的话，请推荐一些适合从Java转到Python的学习资源或练习项目，以辅助他的学习过程。",
    "Python": "你是一位精通Python开发语言的高级工程师，擅长利用多种类库组合来以最简洁有效的方式解决用户提出的需求。此外，你还非常注重代码的质量与结构，并要求对关键部分添加中文注释以便于理解和维护。请根据这一背景信息，帮助我完成我的项目需求。",
}

# 模型选择
chat_model = {
    "lightning": "yi-lightning", # 16K	适用于实时交互，高复杂推理场景，极高的性价比能够为商业产品提供极好的产品支撑	¥0.99
    "turbo": "yi-large-turbo", # 16K	适用于全场景、高品质的推理及文本生成等场景。	¥12
    "large": "yi-large", # 32K	适合于复杂语言理解、深度内容创作设计等复杂场景。	¥20
    "spark": "yi-spark", # 16K	适用于轻量化数学分析、代码生成、文本聊天等场景。	¥1
}

# 初始化最大会话历史长度
if "max_history_length" not in st.session_state:
    st.session_state.max_history_length = 5

if "chat_model" not in st.session_state:
    st.session_state.chat_model = chat_model.get("lightning")

if "message_body" not in st.session_state:
    st.session_state.message_body =[
    {
        "role": "system",
        "content": system_roles.get("Chat Robot")
    }
]

if "history_message" not in st.session_state:
    st.session_state.history_message =[]

# 创建ChatBox对象
chat_box = ChatBox(
    chat_name="default_name",
    session_key="chat_history",
    # user_avatar="👨‍💻",
    # assistant_avatar="🤖",
    greetings=[
        system_roles.get("Chat Robot"),
    ],
)

# 更新最大历史记录长度，如果输入长度小于当前历史记录长度，则重新复制历史消息
def update_max_history_length(input_max_history_length):
    if input_max_history_length < st.session_state.max_history_length:
        new_history_chat = []
        current_history_chat_len = len(st.session_state.history_message)
        for i in range(min(input_max_history_length, current_history_chat_len)):
            new_history_chat.append(st.session_state.history_message[i])
        st.session_state.history_message = new_history_chat
    st.session_state.max_history_length = input_max_history_length



def update_system_role(role_name):
    st.session_state.message_body[0]["content"] = system_roles.get(role_name)

def update_chat_model(module_key):
    st.session_state.chat_model = chat_model.get(module_key)

# 侧边栏配置
with st.sidebar:
    st.subheader('start to chat using streamlit')
    selected_role = st.sidebar.selectbox("选择系统角色", list(system_roles.keys()), index=0)
    update_system_role(selected_role)
    st.divider()
    selected_model = st.sidebar.selectbox("选择模型", list(chat_model.keys()), index=0)
    update_chat_model(selected_model)
    input_max_history_length = st.number_input("设置上下文长度", min_value=1, max_value=10, value=st.session_state.max_history_length)
    update_max_history_length(input_max_history_length)

chat_box.init_session()
chat_box.output_messages()

# 调用大模型并存储历史会话信息
def history_chat(user_input) :

    # 创建一个循环流式获取输出
    while True:
        messages = st.session_state.message_body.copy()
        messages.append({"role": "user", "content": user_input}) # 用户输入内容
        if len(st.session_state.history_message) > 0:
            # 添加历史记录
            for history_chat in st.session_state.history_message:
                messages.append(history_chat)
        # 使用先进先出策略限制历史记录长度
        if len(st.session_state.history_message) >= st.session_state.max_history_length:
            st.session_state.history_message.pop(0)  # 移除最早的记录

        # 使用指定的模型生成回复，并将回复添加到聊天历史记录中
        try:
            response = client.chat.completions.create(
                model=st.session_state.chat_model,
                messages=messages,
                stream=True
            )
            return response
        except Exception as e:
            st.error(f"调用API时出错: {e}")
            return None


# 处理输入信息
if query := st.chat_input('请在这里输入你的问题...'):
    chat_box.user_say(query)
    chat_box.ai_say(
        [
            # you can use string for Markdown output if no other parameters provided
            Markdown("思考中...", expanded=True, title="answer"),
            Markdown("生成信息仅供参考。", title="references"),
        ]
    )

    # 获取模型的回复内容
    model_response = ""
    for chunk in history_chat(query):
        if chunk.choices[0].delta.content is not None:
            model_response += chunk.choices[0].delta.content
            chat_box.update_msg(model_response, element_index=0, streaming=True)

    # 将模型的回复也添加到聊天历史记录中
    st.session_state.history_message.append({"role": "assistant", "content": model_response})
    chat_box.update_msg(model_response, element_index=0, streaming=False, state="complete")
